import brian as brian
from brian import mV, ms, second

    
class SynfireNetwork:
           
    # Neuron model parameters
    Vr = -70*mV
    Vt = -55*mV
    taum = 10*ms
    taupsp = 0.325*ms
    weight = 4.86 * mV
    refractoriness = 1 * ms
    
    # binSize used for packet detection
    binSize = 0.2 * ms 
    Input = None
    SpikeMonitors = None
    
    # class constructor
    def __init__(self, input):
        # create the network (without any inputs)
        self.SetupNetwork()
        self.SetInput(input)
        
    
    # Creates the neurons and their connections, also adds monitors
    def SetupNetwork(self):
        # Neuron model
        eqs = brian.Equations('''
                dV/dt=(-(V-self.Vr)+x)*(1./self.taum) : volt
                dx/dt=(-x+y)*(1./self.taupsp) : volt
                dy/dt=-y*(1./self.taupsp)+25.27*mV/ms+\
                (39.24*mV/ms**0.5)*xi : volt
                ''')
        
        # Neuron groups
        self.Neurons = brian.NeuronGroup(N=1000, model=eqs,
                threshold=self.Vt,reset=self.Vr,refractory= self.refractoriness)
        
        # The network structure
        self.Layers = [self.Neurons.subgroup(100) for i in range(10)]
        self.NeuronConnection = brian.DelayConnection(self.Neurons, self.Neurons,'y', max_delay = 2 * ms)
        
        for i in range(9):
            self.NeuronConnection.connect_full(self.Layers[i], self.Layers[i+1],self.weight, delay = (0 * ms, 2 * ms))
            
        self.Network = brian.Network(self.Neurons, self.NeuronConnection)
        
        self.AddMonitors()
    
    # Extracts the list of spiketimes from the spike monitor
    def GetSpikeTimes(self, spikeMonitor):
        return [spiketime for spiketimes in spikeMonitor.spiketimes.itervalues() for spiketime in spiketimes]
    
    # Calculates the average spike count per  bin (10% are added)
    def CalcAverageSpikeCountPerBin(self):
        self.AverageSpikeCountsPerBin = []
        
        # Create bins and count number of spikes in each bin
        for SpikeMonitor in self.SpikeMonitors:
           averageSpikeCountPerBinOfLayer = SpikeMonitor.nspikes / (self.lastRunTime / self.binSize) * 1.1
           if averageSpikeCountPerBinOfLayer < 2:
               averageSpikeCountPerBinOfLayer = 2
               
           self.AverageSpikeCountsPerBin.append(averageSpikeCountPerBinOfLayer)
                            
      
    # Connects the given input with the first layer of the network
    def SetInput(self, input):
        self.Input = input
        self.InputConnection = brian.Connection(self.Input, self.Layers[0], 'y')
        self.InputConnection.connect_full(weight=self.weight)
        
        self.SpikeMonitors = [brian.SpikeMonitor(self.Input)] + self.SpikeMonitors
        
        self.Network.add(self.Input, self.InputConnection, self.SpikeMonitors)
        
    # Adds all monitors required to the network
    def AddMonitors(self):
        # Setup monitors
        
        # Record state variables like standard deviation
        self.StateMonitor =  [brian.StateMonitor(p, 'y') for p in self.Layers]
        
        # Record the single spikes
        self.SpikeMonitors = [brian.SpikeMonitor(p) for p in self.Layers]
        
        # Record population rates
        self.PSTHMonitor = [brian.PopulationRateMonitor(source,1*ms) for source in self.Layers]
        
        self.Network.add(self.StateMonitor, self.SpikeMonitors, self.PSTHMonitor)
        
        
    # Runs a simulation
    def Run(self, trialTime):
            
        # Setup the network, and run it
        self.Neurons.V = self.Vr + brian.rand(len(self.Neurons)) * (self.Vt-self.Vr)
        
        self.lastRunTime = trialTime
        
        self.Network.run(trialTime)
        self.Network.reinit()
        
        #self.CalcAverageSpikeCountPerBin()

        packet = self.FindPacket(self.SpikeMonitors[-1])
        
        return self.GetClippedStandardDeviation(self.SpikeMonitors[-1], packet)
    
    # Finds a packet within each layer, if available
    def FindPacket(self, SpikeMonitor):
        def GetSpikeCountInBin(SpikeTimes, currentPosition):
            return len(filter(lambda x: x * second >= currentPosition and x * second  < currentPosition + self.binSize, SpikeTimes))
        
        SpikeTimes = self.GetSpikeTimes(SpikeMonitor)
        averageSpikeCountPerBin = 2 #self.AverageSpikeCountsPerBin[self.SpikeMonitors.index(SpikeMonitor) - 1] 
        
        currentPosition = 0 * ms
        
        packetStart = -1 * ms
        packetEnd   = -1 * ms
        searchEnd   = False
        
        while currentPosition < self.lastRunTime:
            if GetSpikeCountInBin(SpikeTimes, currentPosition) > averageSpikeCountPerBin:
                if packetStart == -1 * ms:
                    packetStart = currentPosition
                
                searchEnd = True
            else:
                if searchEnd:
                    packetEnd = currentPosition
                    searchEnd = False
                
               
            currentPosition += self.binSize
        
        if packetStart != -1 * ms:
            if searchEnd:
                packetEnd = self.lastRunTime
             
            return (packetStart, packetEnd)
        else:
            return None
    
    # Returns the standard deviation of all spike within a packet detected
    # or the standard deviation of all spikes if no packet has been found
    def GetClippedStandardDeviation(self, SpikeMonitor, packetPosition):
        SpikeTimes = self.GetSpikeTimes(SpikeMonitor)
        
        if packetPosition is None:
            return brian.std(SpikeTimes)
        else:
            return brian.std(filter(lambda x: x * second >= packetPosition[0]  and x * second <= packetPosition[1], SpikeTimes))
    
    # Plots the results of the last test run
    def ShowResults(self):
        
        # Plot results
        
        # Plot spikes by time
        brian.raster_plot(showgrouplines=True,*self.SpikeMonitors)
        
        for packet in self.Packets:
            if not packet is None:
                index = self.Packets.index(packet)
                brian.plot([packet[0] / ms, packet[0] / ms], [index, index + 1], 'g--', linewidth=2)
                brian.plot([packet[1] / ms, packet[1] / ms], [index, index + 1], 'r--', linewidth=2)
        
            
        brian.xlabel("time (ms)")
        brian.xlabel("population rate (Hz)")
            
        
        brian.show()
        
    


spikeCounts = range(10, 100, 10)
sigmas = range(0, 5) 

brian.subplot(211)
brian.xlabel(r'$\sigma$')
brian.ylabel(r'$n_{spikes}$')


for n in spikeCounts:
    for sigma in sigmas:
        model = SynfireNetwork(brian.PulsePacket(t=10*ms,n=n,sigma=sigma*ms))
        std = model.Run(70 * ms)
        
        if std < 0.01:
            marker = 'x'
        else:
            marker = 'o'
        
        # Plot result
        # Determine color (r, g, b), the smaller the standard deviation, the greener the result
        color = (min(0.7 * std * 50, 0.7), max(0.7 * (1 - (std * 50)), 0), 0)
        brian.plot(sigma, n, color = color, marker = 'x', markeredgewidth=3, markersize=12)
        

brian.xlim(sigmas[0] - 1, sigmas[-1] + 1)
brian.ylim(spikeCounts[0] - 5, spikeCounts[-1] + 5)

 # Plot spikes by time
brian.subplot(212)
brian.raster_plot(showgrouplines=True,*model.SpikeMonitors)    
brian.xlabel("time (ms)")
brian.xlabel("population rate (Hz)")

brian.show()
